import pandas as pd
import os
import re
import logging
import pypinyin  # 新增导入拼音库


class FuzzySearch:
    def __init__(self, db_path='resources/stock_database.csv'):
        self.db_path = db_path
        self.stock_db = []
        self.load_database()

    def load_database(self):
        """加载股票数据库"""
        try:
            if os.path.exists(self.db_path):
                # 尝试使用不同编码读取CSV文件 - 优先尝试中文编码
                encodings = ['gbk', 'gb2312', 'utf-8', 'latin1']
                df = None

                for encoding in encodings:
                    try:
                        df = pd.read_csv(self.db_path, encoding=encoding)
                        # 检查列名是否正确
                        if 'symbol' in df.columns and 'name' in df.columns and 'market' in df.columns:
                            logging.info(f"使用 {encoding} 编码加载股票数据库成功")
                            break
                    except Exception as e:
                        logging.warning(f"使用 {encoding} 编码加载失败: {str(e)}")
                        continue

                # 如果未成功读取，尝试使用第一行作为列名
                if df is None:
                    logging.warning("尝试使用第一行作为列名")
                    for encoding in encodings:
                        try:
                            df = pd.read_csv(self.db_path, encoding=encoding, header=0)
                            if 'symbol' in df.columns and 'name' in df.columns and 'market' in df.columns:
                                logging.info(f"使用 {encoding} 编码加载股票数据库成功（使用第一行作为列名）")
                                break
                        except:
                            continue

                if df is not None and not df.empty:
                    # 确保有必要的列
                    required_columns = ['symbol', 'name', 'market']
                    if all(col in df.columns for col in required_columns):
                        self.stock_db = [
                            (str(row['symbol']), str(row['name']), str(row['market']))
                            for _, row in df.iterrows()
                        ]
                        logging.info(f"加载股票数据库成功，共 {len(self.stock_db)} 条记录")
                    else:
                        logging.error(f"CSV文件缺少必要列: {required_columns}")
                        self.stock_db = []
                else:
                    logging.error(f"无法加载股票数据库文件: {self.db_path}")
                    self.stock_db = []
            else:
                logging.warning(f"股票数据库文件不存在: {self.db_path}")
        except Exception as e:
            logging.error(f"加载股票数据库失败: {str(e)}")
            self.stock_db = []

    def search(self, query, max_results=10):
        """执行模糊搜索"""
        try:
            if not self.stock_db or not query:
                return []

            # 清理查询：去除特殊字符
            cleaned_query = re.sub(r'[^\w\u4e00-\u9fa5]', '', query).lower()

            # 如果清理后为空，直接返回
            if not cleaned_query:
                return []

            results = []

            # 如果查询是纯数字，尝试匹配股票代码
            if cleaned_query.isdigit():
                # 匹配以查询开头的股票代码
                results = [
                    (symbol, name, market)
                    for symbol, name, market in self.stock_db
                    if symbol.startswith(cleaned_query)
                ]

                # 如果结果不足，尝试包含匹配
                if len(results) < max_results:
                    results += [
                                   (symbol, name, market)
                                   for symbol, name, market in self.stock_db
                                   if cleaned_query in symbol and (symbol, name, market) not in results
                               ][:max_results - len(results)]

            # 尝试匹配股票名称
            if not results:
                # 名称完全匹配
                results = [
                    (symbol, name, market)
                    for symbol, name, market in self.stock_db
                    if cleaned_query == name.lower()
                ]

                # 名称包含匹配
                if len(results) < max_results:
                    results += [
                                   (symbol, name, market)
                                   for symbol, name, market in self.stock_db
                                   if cleaned_query in name.lower() and (symbol, name, market) not in results
                               ][:max_results - len(results)]

                # 拼音首字母匹配 - 使用pypinyin库
                if len(results) < max_results:
                    pinyin_initials = self.get_pinyin_initials(cleaned_query)
                    if pinyin_initials:
                        results += [
                                       (symbol, name, market)
                                       for symbol, name, market in self.stock_db
                                       if self.get_pinyin_initials(name).startswith(pinyin_initials) and
                                          (symbol, name, market) not in results
                                   ][:max_results - len(results)]

            return results[:max_results]
        except Exception as e:
            logging.error(f"执行模糊搜索时出错: {str(e)}")
            return []

    def get_pinyin_initials(self, text):
        """获取中文文本的拼音首字母 - 使用pypinyin库"""
        if not text:  # 处理空文本
            return ""

        try:
            # 使用pypinyin获取拼音首字母
            initials = ''.join([p[0] for p in pypinyin.lazy_pinyin(text, style=pypinyin.Style.FIRST_LETTER)])
            return initials.lower()
        except Exception as e:
            logging.error(f"获取拼音首字母时出错: {str(e)}")
            return ""